from torchvision.models.resnet import resnet50
from torchvision import models
import torch.nn as nn
import torch
import torch.nn.functional as F


class NeuralNetwork(nn.Module):
    def __init__(self, input_feature, num_classes):
        super().__init__()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(input_feature, 128),
            # nn.ReLU(),
            # nn.Linear(512, 512),
            nn.ReLU(),
            nn.Linear(128, num_classes),
        )

    def forward(self, x):
        logits = self.linear_relu_stack(x)
        pred_probab = nn.LogSoftmax(dim=1)(logits)
        return pred_probab

class NN_4layer(nn.Module):
    def __init__(self, input_feature, num_classes):
        super().__init__()
        self.fc1 = nn.Linear(input_feature, 1024)
        self.fc2 = nn.Linear(1024, 512)
        self.fc3 = nn.Linear(512, 256)
        self.fc4 = nn.Linear(256, num_classes)
        self.Tanh = nn.Tanh()

    def forward(self, x):
        x = self.Tanh(self.fc1(x))
        x = self.Tanh(self.fc2(x))
        x = self.Tanh(self.fc3(x))
        logits = self.fc4(x)
        # pred_probab = nn.LogSoftmax(dim=1)(logits)
        return logits

class M_CNN(nn.Module):
    def __init__(self, num_classes):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 16, 8, 2, padding=3)
        self.conv2 = nn.Conv2d(16, 32, 4, 2)
        self.fc1 = nn.Linear(32 * 4 * 4, 32)
        self.fc2 = nn.Linear(32, num_classes)

    def forward(self, x):
        # x of shape [B, 1, 28, 28]
        x = F.relu(self.conv1(x))  # -> [B, 16, 14, 14]
        x = F.max_pool2d(x, 2, 1)  # -> [B, 16, 13, 13]
        x = F.relu(self.conv2(x))  # -> [B, 32, 5, 5]
        x = F.max_pool2d(x, 2, 1)  # -> [B, 32, 4, 4]
        x = x.view(-1, 32 * 4 * 4)  # -> [B, 512]
        x = F.relu(self.fc1(x))  # -> [B, 32]
        x = self.fc2(x)  # -> [B, 10]
        # pred_probab = nn.LogSoftmax(dim=1)(x)
        return x

class CNN(nn.Module):
    def __init__(self, num_classes):
        super().__init__()
        self.conv1 = nn.Conv2d(3,16,3,padding=1)
        self.conv2 = nn.Conv2d(16,32,3,padding=1)
        self.conv3 = nn.Conv2d(32,64,3,padding=1)
        self.pool = nn.MaxPool2d(2,2)
        self.fc1 = nn.Linear(64*4*4,500)
        self.fc2 = nn.Linear(500,num_classes)
        self.BatchNorm1 = nn.BatchNorm2d(16)
        self.BatchNorm2 = nn.BatchNorm2d(32)
        self.BatchNorm3 = nn.BatchNorm2d(64)

    def forward(self, x):
        x = self.pool(F.relu(self.BatchNorm1(self.conv1(x))))
        x = self.pool(F.relu(self.BatchNorm2(self.conv2(x))))
        x = self.pool(F.relu(self.BatchNorm3(self.conv3(x))))
        x = x.view(-1,64*4*4)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        # pred_probab = nn.LogSoftmax(dim=1)(x)
        return x

def create_ResNet50_model(num_classes):
    model = resnet50(weights=None)
    fc_inputs = model.fc.in_features
    model.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=(2, 2), padding=1, bias=False) 
    model.fc = nn.Sequential(         
        nn.Linear(fc_inputs, num_classes),  
        nn.LogSoftmax(dim=1)
    )
    return model

def create_ResNet18_model(num_classes):
    model = ResNet(ResidualBlock, num_classes)
    # model = ResNet_18(ResidualBlock, num_classes)
    # model = models.resnet18(num_classes)
    return model


def create_CNN_model(num_classes, data_name):
    if data_name == 'MNIST':
        model = M_CNN(num_classes)
    else:
        model = CNN(num_classes)
    return model

def create_NN_model(feature, num_classes):
    model = NeuralNetwork(feature, num_classes)
    return model

def create_NN_4layer_model(feature, num_classes):
    model = NN_4layer(feature, num_classes)
    return model


def conv_shortcut(in_channel,out_channel,stride):
    layers = [nn.Conv2d(in_channel,out_channel,kernel_size=(1,1),stride=(stride,stride)),
             nn.BatchNorm2d(out_channel)]
    return nn.Sequential(*layers)

def block(in_channel,out_channel,k_size,stride, conv=False):
    layers = None
    
    first_layers = [nn.Conv2d(in_channel,out_channel[0],kernel_size=(1,1),stride=(1,1)),
                    nn.BatchNorm2d(out_channel[0]),
                    nn.ReLU(inplace=True)]
    if conv:
        first_layers[0].stride=(stride,stride)
    
    second_layers = [nn.Conv2d(out_channel[0],out_channel[1],kernel_size=(k_size,k_size),stride=(1,1),padding=1),
                    nn.BatchNorm2d(out_channel[1])]

    layers = first_layers + second_layers
    
    return nn.Sequential(*layers)
    

class MResnet(nn.Module):
    
    def __init__(self,in_channels,num_classes):
        super().__init__()
        
        self.stg1 = nn.Sequential(
                                   nn.Conv2d(in_channels=in_channels,out_channels=64,kernel_size=(3),stride=(1),padding=1),
                                   nn.BatchNorm2d(64),
                                   nn.ReLU(inplace=True),
                                   nn.MaxPool2d(kernel_size=3,stride=2))
        
        ##stage 2
        self.convShortcut2 = conv_shortcut(64,256,1)
        
        self.conv2 = block(64,[64,256],3,1,conv=True)
        self.ident2 = block(256,[64,256],3,1)

        
        ##stage 3
        self.convShortcut3 = conv_shortcut(256,512,2)
        
        self.conv3 = block(256,[128,512],3,2,conv=True)
        self.ident3 = block(512,[128,512],3,2)

        
        ##stage 4
        self.convShortcut4 = conv_shortcut(512,1024,2)
        
        self.conv4 = block(512,[256,1024],3,2,conv=True)
        self.ident4 = block(1024,[256,1024],3,2)
        
        
        ##Classify
        self.classifier = nn.Sequential(
                                       nn.AvgPool2d(kernel_size=(4)),
                                       nn.Flatten(),
                                       nn.Linear(1024,num_classes))
        
    def forward(self,inputs):
        out = self.stg1(inputs)
        
        #stage 2
        out = F.relu(self.conv2(out) + self.convShortcut2(out))
        out = F.relu(self.ident2(out) + out)
        out = F.relu(self.ident2(out) + out)
        
        #stage3
        out = F.relu(self.conv3(out) + (self.convShortcut3(out)))
        out = F.relu(self.ident3(out) + out)
        out = F.relu(self.ident3(out) + out)
        out = F.relu(self.ident3(out) + out)
        
#         stage4             
        out = F.relu(self.conv4(out) + (self.convShortcut4(out)))
        out = F.relu(self.ident4(out) + out)
        out = F.relu(self.ident4(out) + out)
        out = F.relu(self.ident4(out) + out)
        out = F.relu(self.ident4(out) + out)
        out = F.relu(self.ident4(out) + out)
        
        out = self.classifier(out)#100x1024
        
        return out
    

def create_MResnet_model(num_classes):
    model = MResnet(3, num_classes)
    return model


class ResidualBlock(nn.Module):
    def __init__(self, inchannel, outchannel, stride=1):
        super(ResidualBlock, self).__init__()
        self.left = nn.Sequential(
            nn.Conv2d(inchannel, outchannel, kernel_size=3, stride=stride, padding=1, bias=False),
            nn.BatchNorm2d(outchannel),
            nn.ReLU(inplace=True),
            nn.Conv2d(outchannel, outchannel, kernel_size=3, stride=1, padding=1, bias=False),
            nn.BatchNorm2d(outchannel)
        )
        self.shortcut = nn.Sequential()
        if stride != 1 or inchannel != outchannel:
            self.shortcut = nn.Sequential(
                nn.Conv2d(inchannel, outchannel, kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(outchannel)
            )

    def forward(self, x):
        out = self.left(x)
        out += self.shortcut(x)
        out = F.relu(out)
        return out

class ResNet(nn.Module):
    def __init__(self, ResidualBlock, num_classes=10):
        super(ResNet, self).__init__()
        self.inchannel = 64
        self.conv1 = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False),
            nn.BatchNorm2d(64),
            nn.ReLU(),
        )
        self.layer1 = self.make_layer(ResidualBlock, 64,  2, stride=1)
        self.layer2 = self.make_layer(ResidualBlock, 128, 2, stride=2)
        self.layer3 = self.make_layer(ResidualBlock, 256, 2, stride=2)
        self.layer4 = self.make_layer(ResidualBlock, 512, 2, stride=2)
        self.fc = nn.Linear(512, num_classes)

    def make_layer(self, block, channels, num_blocks, stride):
        strides = [stride] + [1] * (num_blocks - 1)   #strides=[1,1]
        layers = []
        for stride in strides:
            layers.append(block(self.inchannel, channels, stride))
            self.inchannel = channels
        return nn.Sequential(*layers)

    def forward(self, x):
        out = self.conv1(x)
        out = self.layer1(out)
        out = self.layer2(out)
        out = self.layer3(out)
        out = self.layer4(out)
        out = F.avg_pool2d(out, 4)
        out = out.view(out.size(0), -1)
        out = self.fc(out)
        return out

class ResNet_18(nn.Module):
    def __init__(self,ResBlock,num_classes):
        super(ResNet_18, self).__init__()
 
        self.in_channels = 64 #输入layer1时的channel
        #第一层单独卷积层
        self.conv1=nn.Sequential(
            # (n-f+2*p)/s+1,n=28,n=32
            # nn.Conv2d(in_channels=3,out_channels=64,kernel_size=7,stride=2,padding=3),
            nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3, stride=1, padding=1), #64
            nn.BatchNorm2d(64),
            nn.ReLU(),
            # nn.MaxPool2d(kernel_size=3,stride=2,padding=1)
            nn.MaxPool2d(kernel_size=1, stride=1, padding=0) #64
            # nn.Dropout(0.25)
        )
 
        self.layer1=self.make_layer(ResBlock,64,2,stride=1) #64
        self.layer2 = self.make_layer(ResBlock, 128, 2, stride=2) #32
        self.layer3 = self.make_layer(ResBlock, 256, 2, stride=2) #16
        self.layer4 = self.make_layer(ResBlock, 512, 2, stride=2) #8
 
 
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) #torch.nn.AdaptiveAvgPool2d()接受两个参数，分别为输出特征图的长和宽，其通道数前后不发生变化。
                                                    #即这里将输入图片像素强制转换为1*1
        # self.linear=nn.Linear(2*2*512,512)
        # self.linear2=nn.Linear(512,100)
 
        self.linear=nn.Linear(512*1*1,num_classes)
 
        self.dropout = nn.Dropout(0.3)
 
    # 这个函数主要是用来，重复同一个残差块
    def make_layer(self, block, out_channels, num_blocks, stride):
        strides = [stride] + [1] * (num_blocks - 1)
        layers = []
        for stride in strides:
            layers.append(block(self.in_channels, out_channels, stride))
            self.in_channels = out_channels
        return nn.Sequential(*layers)
 
    def forward(self, x):
        x=self.conv1(x)
        # x=self.dropout(x)
        x=self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        x=self.avgpool(x)
        x = x.view(x.size(0), -1)
        x=self.linear(x)
        x=self.dropout(x)
 
        return x








